import pandas as pd

from app_config import get_engine_ts
from app_config import get_pro
import os

from datetime import datetime
from dateutil.relativedelta import relativedelta


class O000300(object):
    one_year_ago: str = None

    zero_year_ago: str = None

    folder_prefix: str = None

    date_000300_for_remain: str = None

    date_000300_for_merge: str = None

    zzqz_date: str = None


""" 沪深300
"""


def calc_hs300(o_000300: O000300 | None = None):
    one_year_ago = o_000300.one_year_ago
    zero_year_ago = o_000300.zero_year_ago
    str_file_pre = o_000300.folder_prefix
    date_000300_for_merge = o_000300.date_000300_for_merge
    date_000300_for_remain = o_000300.date_000300_for_remain
    zzqz_date = o_000300.zzqz_date

    print(f"""
      one_year_ago: {"   " + one_year_ago}
      zero_year_ago: {"  " + zero_year_ago}
      date_000300_for_remain: {" " + date_000300_for_remain}
      date_000300_for_merge: {" " + date_000300_for_merge}
     zzqz_date: {" " + zzqz_date}
     str_file_pre: {"    " + str_file_pre}
 
      """)

    if not os.path.exists(str_file_pre):
        os.makedirs(str_file_pre)
        print(f"文件夹 '{str_file_pre}' 创建成功")
    engine = get_engine_ts()
    """     1 样本空间    """
    ''' 
        指数样本空间由同时满足以下条件的非 ST、*ST 沪深 A 股和红筹企业发行的存托凭证组成：
             科创板证券、创业板证券：上市时间超过一年。
             其他证券：上市时间超过一个季度，除非该证券自上市以来日均总市值排在前 30 位
    '''
    # 同中证全指
    df_stock_basic = get_pro().index_weight(index_code='000985.CSI', start_date=zzqz_date, end_date=zzqz_date)
    df_stock_basic.to_excel(str_file_pre + '中证全指.xlsx')
    list_index_ts_code = df_stock_basic['con_code'].tolist()
    query_hfd = f"""
           SELECT ts_code, symbol, name, market, industry, list_date FROM `stock_basic`
           WHERE ts_code IN ({','.join(f"'{code}'" for code in list_index_ts_code)})
           AND market != '北交所'
           """
    df_stock_basic = pd.read_sql_query(query_hfd, engine)

    print("=====" + str(len(df_stock_basic)))

    df_stock_basic = df_stock_basic[
        ~((df_stock_basic['market'] == '科创板') & (df_stock_basic['list_date'] > one_year_ago))]
    print("=====" + str(len(df_stock_basic)))

    '''
         对样本空间内证券按照过去一年的日均成交金额由高到低排名，剔除排名后 50%的证券；
         对样本空间内剩余证券，按照过去一年的日均总市值由高到低排名，选取前 300 名的证券作为指数样本。
    '''
    filter_ts_code = df_stock_basic['ts_code'].tolist()
    # 构建查询语句
    query = f"""
    SELECT ts_code,trade_date, amount FROM `daily`
    WHERE ts_code IN ({','.join(f"'{code}'" for code in filter_ts_code)})
    AND trade_date >= '{one_year_ago}'
    AND trade_date <= '{zero_year_ago}'
    """

    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)
    result_df['amount'] = result_df['amount'] / 1000
    grouped_df = result_df.groupby('ts_code')['amount'].mean().reset_index()
    grouped_df.columns = ['ts_code', '日均成交额']
    merge = pd.merge(df_stock_basic, grouped_df, on='ts_code', how='left')

    # 构建查询语句
    query = f"""
        SELECT ts_code,trade_date, total_mv FROM `daily_basic`
        WHERE ts_code IN ({','.join(f"'{code}'" for code in filter_ts_code)})
        AND trade_date >= '{one_year_ago}'
        AND trade_date <= '{zero_year_ago}'
        """
    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)
    result_df['total_mv'] = result_df['total_mv'] / 10000
    grouped_df_total_mv = result_df.groupby('ts_code')['total_mv'].mean().reset_index()
    grouped_df_total_mv.columns = ['ts_code', '日均总市值']
    pd_merge = pd.merge(merge, grouped_df_total_mv, on='ts_code', how='left')
    pd_merge.to_excel(str_file_pre + '日均成交额.xlsx')

    pre_000300_index = get_pro().index_weight(index_code='000300.SH', start_date=date_000300_for_remain,
                                              end_date=date_000300_for_remain)
    pre_000300_index.to_excel(str_file_pre + '沪深300_上期.xlsx', index=False)
    # pre_000300_index.rename(columns={'con_code': 'ts_code'}, inplace=True)
    my_set = pre_000300_index['con_code'].tolist()

    # 按日均成交额 列降序排序
    df_sorted = pd_merge.sort_values(by='日均成交额', ascending=False).reset_index(drop=True)
    # 获取前 90% 的数据
    num_rows = len(df_sorted)
    # df_sorted['i_avgAmount'] = df_sorted.index.astype(str) + '/' + str(int(num_rows * 0.5))
    df_sorted['i_日均成交'] = df_sorted.index.astype(int) + 1
    df_sorted.loc[df_sorted['ts_code'].isin(my_set), 'i_日均成交'] = 0
    df_sorted = df_sorted.sort_values(by='i_日均成交').reset_index(drop=True)
    top_50_percent = df_sorted.head(int(num_rows * 0.5))

    top_50_percent = top_50_percent.sort_values('日均总市值', ascending=False).reset_index(drop=True)
    top_50_percent['i_日均市值'] = top_50_percent.index
    top_50_percent.to_excel(str_file_pre + '沪深300_top50.xlsx')
    df_result_hs300 = top_50_percent.head(400)
    # print("开始合并")
    df_000300 = get_pro().index_weight(index_code='000300.SH', start_date=date_000300_for_merge, end_date=date_000300_for_merge)
    df_000300.rename(columns={'con_code': 'ts_code'}, inplace=True)
    df_000300.rename(columns={'trade_date': 'list_date'}, inplace=True)
    df_000300['symbol'] = df_000300['ts_code'].str.split('.').str[0]
    df_000300.drop(columns=['weight'], inplace=True)

    df_000300 = pd.merge(df_000300, df_sorted[['ts_code', 'i_日均成交', '日均成交额']], on='ts_code', how='left')
    df_000300.to_excel(str_file_pre + date_000300_for_merge + '-000300.xlsx')

    # 获取 dataframe1 中的所有列
    columns_dataframe1 = df_result_hs300.columns

    # 创建缺失列并填充为 '*'
    for column in columns_dataframe1:
        if column not in df_000300.columns:
            df_000300[column] = '**'

    # 将 dataframe2 的列顺序调整为与 dataframe1 一致
    df_000300 = df_000300[columns_dataframe1]

    # 追加 dataframe2 到 dataframe1
    result = pd.concat([df_result_hs300, df_000300], ignore_index=True)

    result.to_excel(str_file_pre + '沪深300.xlsx')


# strStartDate: 周期开启时间
# strEndDate: 周期结束时间
# str_filePre: 输出结果到文件夹
# pre_000300: 上一期沪深300成分股
# cur_000300: 本期沪深300成分股
if __name__ == '__main__':
    # """2024年4季度"""
    strStartDate = '20231101'
    strEndDate = '20241031'
    folder_prefix = 'zfile/沪深300_2024季度4/'
    pre_000300 = '20240930'  # 上一期沪深300 将上一期的所有股票进入日均总市值样本
    cur_000300 = '20241231'  # 本期沪深300
    calc_hs300(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=folder_prefix, cur_000300Date=cur_000300,
               pre_000300Date=pre_000300)

    # """2024年2季度"""
    strStartDate = '20230501'
    strEndDate = '20240430'
    folder_prefix = 'zfile/沪深300_2024季度2/'
    pre_000300 = '20240329'  # 上一期沪深300 将上一期的所有股票进入日均总市值样本
    cur_000300 = '20240628'  # 本期沪深300
    calc_hs300(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=folder_prefix, cur_000300Date=cur_000300,
               pre_000300Date=pre_000300)

    # """2023年四季度"""
    strStartDate = '20221101'
    strEndDate = '20231031'
    folder_prefix = 'zfile/沪深300_2023季度4/'
    pre_000300 = '20230928'  # 上一期沪深300
    cur_000300 = '20231229'
    calc_hs300(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=folder_prefix, cur_000300Date=cur_000300,
               pre_000300Date=pre_000300)

    # """2023年二季度"""
    strStartDate = '20220501'
    strEndDate = '20230430'
    folder_prefix = 'zfile/沪深300_2023季度2/'
    pre_000300 = '20230331'
    cur_000300 = '20230630'
    calc_hs300(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=folder_prefix, cur_000300Date=cur_000300,
               pre_000300Date=pre_000300)

    # """2022年四季度"""
    strStartDate = '20211101'
    strEndDate = '20221031'
    folder_prefix = 'zfile/沪深300_2022季度4/'
    pre_000300 = '20220930'
    cur_000300 = '20221230'
    calc_hs300(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=folder_prefix, cur_000300Date=cur_000300,
               pre_000300Date=pre_000300)

    # """2022年二季度"""
    strStartDate = '20210501'
    strEndDate = '20220429'
    folder_prefix = 'zfile/沪深300_2022季度2/'
    pre_000300 = '20220331'
    cur_000300 = '20220630'
    calc_hs300(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=folder_prefix, cur_000300Date=cur_000300,
               pre_000300Date=pre_000300)

    # """2021年四季度"""
    strStartDate = '20201101'
    strEndDate = '20211029'
    folder_prefix = 'zfile/沪深300_2021季度4/'
    pre_000300 = '20210930'
    cur_000300 = '20211231'
    calc_hs300(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=folder_prefix, cur_000300Date=cur_000300,
               pre_000300Date=pre_000300)

    # """2021年二季度"""
    strStartDate = '20200501'
    strEndDate = '20210430'
    folder_prefix = 'zfile/沪深300_2021季度2/'
    pre_000300 = '20210331'
    cur_000300 = '20210630'
    calc_hs300(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=folder_prefix, cur_000300Date=cur_000300,
               pre_000300Date=pre_000300)
